package aes_go

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"io"
	"log"
)

type AesG struct {
	key []byte
}

func checKeyLen(num int, key []byte) []byte {
	for i := 0; i < num; i++ {
		key = append(key, byte('x'))
	}

	return key
}

func NewAes(key []byte) *AesG {

	keyLen := len(key)

	switch keyLen {
	case 16, 24, 32:
		return &AesG{
			key: key,
		}
	}

	if keyLen < 16 {
		appenNumber := 16 - keyLen
		key = checKeyLen(appenNumber, key)
	} else if keyLen < 24 {
		appenNumber := 24 - keyLen
		key = checKeyLen(appenNumber, key)
	} else if keyLen < 32 {
		appenNumber := 32 - keyLen
		key = checKeyLen(appenNumber, key)
	} else if keyLen > 32 {
		panic("aes key must lt 32")
	}
	return &AesG{
		key: key,
	}

}

// Encode 加密
func (a *AesG) Encode(data []byte) (string, error) {
	block, err := aes.NewCipher(a.key)
	if err != nil {
		return "", err
	}
	encrypted := make([]byte, aes.BlockSize+len(data))
	iv := encrypted[:aes.BlockSize]
	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
		return "", err
	}
	stream := cipher.NewCFBEncrypter(block, iv)
	stream.XORKeyStream(encrypted[aes.BlockSize:], data)
	return fmt.Sprintf("%x", encrypted), nil
}

// Decode 解密
func (a *AesG) Decode(encrypted string) (string, error) {

	codePass, err := hex.DecodeString(encrypted)
	if err != nil {
		log.Println("hex DecodeString Err:", err)
		return "", err
	}

	block, _ := aes.NewCipher(a.key)
	if len(encrypted) < aes.BlockSize {
		return "", err
	}
	iv := codePass[:aes.BlockSize]
	codePass = codePass[aes.BlockSize:]

	stream := cipher.NewCFBDecrypter(block, iv)
	stream.XORKeyStream(codePass, codePass)
	return string(codePass), nil
}
